package ace.module.oauth2.server.core.impl.authentication.Introspection;

import ace.cmp.core.exception.BusinessException;
import ace.cmp.security.core.impl.constant.Oauth2Constant;
import ace.cmp.security.core.impl.model.dto.Oauth2UserInfoDto;
import ace.module.oauth2.server.api.callback.Oauth2AceAuthorizationCallbackApi;
import ace.module.oauth2.server.api.callback.model.input.Oauth2AceAuthorizationCallbackInput;
import ace.module.oauth2.server.api.model.dto.ext.Oauth2AceUserInfoDto;
import ace.module.oauth2.server.api.model.dto.ext.Oauth2AceUserInfoResultDto;
import ace.module.oauth2.server.core.impl.authentication.ace.factory.CallbackApiFactory;
import ace.module.oauth2.server.core.impl.dao.entity.Oauth2RegisteredClientAuthConfig;
import ace.module.oauth2.server.core.impl.manager.Oauth2RegisteredClientAuthConfigManager;
import java.net.URL;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.core.convert.TypeDescriptor;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2ErrorCodes;
import org.springframework.security.oauth2.core.OAuth2Token;
import org.springframework.security.oauth2.core.OAuth2TokenIntrospectionClaimNames;
import org.springframework.security.oauth2.core.converter.ClaimConversionService;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenIntrospection;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2TokenIntrospectionAuthenticationProvider;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2TokenIntrospectionAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.stereotype.Component;

/**
 * 扩展 {@link OAuth2TokenIntrospectionAuthenticationProvider}
 * 1. 添加动态获取用户权限.
 *
 * @author caspar
 */
@Slf4j
@AllArgsConstructor
@Component
public class IntrospectionAuthenticationProvider implements AuthenticationProvider {

  private static final TypeDescriptor OBJECT_TYPE_DESCRIPTOR = TypeDescriptor.valueOf(Object.class);
  private static final TypeDescriptor LIST_STRING_TYPE_DESCRIPTOR =
      TypeDescriptor.collection(List.class, TypeDescriptor.valueOf(String.class));
  private final OAuth2AuthorizationService authorizationService;
  private final Oauth2RegisteredClientAuthConfigManager oauth2RegisteredClientAuthConfigManager;
  private final CallbackApiFactory callbackApiFactory;


  @Override
  public Authentication authenticate(Authentication authentication) throws AuthenticationException {
    OAuth2TokenIntrospectionAuthenticationToken tokenIntrospectionAuthentication =
        (OAuth2TokenIntrospectionAuthenticationToken) authentication;

    OAuth2ClientAuthenticationToken clientPrincipal =
        getAuthenticatedClientElseThrowInvalidClient(tokenIntrospectionAuthentication);

    OAuth2Authorization authorization = this.authorizationService.findByToken(tokenIntrospectionAuthentication.getToken(), null);
    if (authorization == null) {
      if (this.log.isTraceEnabled()) {
        this.log.trace(
            "Did not authenticate token introspection input since token was not found");
      }
      // Return the authentication input when token not found
      return tokenIntrospectionAuthentication;
    }

    if (this.log.isTraceEnabled()) {
      this.log.trace("Retrieved authorization with token");
    }

    OAuth2Authorization.Token<OAuth2Token> authorizedToken =
        authorization.getToken(tokenIntrospectionAuthentication.getToken());
    if (!authorizedToken.isActive()) {
      if (this.log.isTraceEnabled()) {
        this.log.trace("Did not introspect token since not active");
      }
      return new OAuth2TokenIntrospectionAuthenticationToken(
          tokenIntrospectionAuthentication.getToken(),
          clientPrincipal,
          OAuth2TokenIntrospection.builder().build());
    }

    RegisteredClient authorizedClient = clientPrincipal.getRegisteredClient();

    OAuth2TokenIntrospection tokenClaims =
        withActiveTokenClaims(authorizedToken, authorizedClient, authorization);

    if (this.log.isTraceEnabled()) {
      this.log.trace("Authenticated token introspection input");
    }

    return new OAuth2TokenIntrospectionAuthenticationToken(
        authorizedToken.getToken().getTokenValue(), clientPrincipal, tokenClaims);
  }

  @Override
  public boolean supports(Class<?> authentication) {
    return OAuth2TokenIntrospectionAuthenticationToken.class.isAssignableFrom(authentication);
  }

  private OAuth2TokenIntrospection withActiveTokenClaims(
      OAuth2Authorization.Token<OAuth2Token> authorizedToken,
      RegisteredClient authorizedClient,
      OAuth2Authorization authorization) {

    OAuth2TokenIntrospection.Builder tokenClaims;
    if (!MapUtils.isEmpty(authorizedToken.getClaims())) {
      Map<String, Object> claims = convertClaimsIfNecessary(authorizedToken.getClaims());
      tokenClaims = OAuth2TokenIntrospection.withClaims(claims).active(true);
    } else {
      tokenClaims = OAuth2TokenIntrospection.builder(true);
    }

    tokenClaims.clientId(authorizedClient.getClientId());

    // TODO Set "username"

    OAuth2Token token = authorizedToken.getToken();
    if (token.getIssuedAt() != null) {
      tokenClaims.issuedAt(token.getIssuedAt());
    }
    if (token.getExpiresAt() != null) {
      tokenClaims.expiresAt(token.getExpiresAt());
    }

    if (OAuth2AccessToken.class.isAssignableFrom(token.getClass())) {
      OAuth2AccessToken accessToken = (OAuth2AccessToken) token;
      tokenClaims.tokenType(accessToken.getTokenType().getValue());
    }

    tokenClaims = this.customizerTokenClaims(authorizedToken, authorizedClient, authorization, tokenClaims);

    return tokenClaims.build();
  }

  /**
   * 自定义token内容
   * 添加用户权限信息
   *
   * @param authorizedToken
   * @param authorizedClient
   * @param tokenClaims
   */
  private OAuth2TokenIntrospection.Builder customizerTokenClaims(
      OAuth2Authorization.Token<OAuth2Token> authorizedToken,
      RegisteredClient authorizedClient,
      OAuth2Authorization authorization,
      OAuth2TokenIntrospection.Builder tokenClaims) {
    if (!this.containUserInfoScope(authorization.getAuthorizedScopes())) {
      return tokenClaims;
    }

    String accountId = authorization.getPrincipalName();
    String idOfClient = authorization.getRegisteredClientId();

    List<Oauth2RegisteredClientAuthConfig> oauth2RegisteredClientAuthorizationConfigs = oauth2RegisteredClientAuthConfigManager
        .getByIdOfClientAndStatusIsEnableAndBizAuthTypeIsAuthorizationTypeFromDbOrCache(idOfClient);

    if (CollectionUtils.isEmpty(oauth2RegisteredClientAuthorizationConfigs)) {
      Oauth2UserInfoDto userInfoResultDetailDto = Oauth2UserInfoDto.builder()
          .id(accountId)
          .attributes(new HashMap<>())
          .authorities(new HashSet<>())
          .build();
      tokenClaims.claim(Oauth2Constant.TOKEN_CLAIMS_USERINFO_ATTRIBUTE_KEY, userInfoResultDetailDto);
      return tokenClaims;
    }

    Oauth2AceAuthorizationCallbackInput userInfoCallbackInput = Oauth2AceAuthorizationCallbackInput.builder()
        .accountId(accountId)
        .idOfClient(idOfClient)
        .build();

    Oauth2AceUserInfoResultDto oauth2AceUserInfoResultDto = null;
    for (Oauth2RegisteredClientAuthConfig oauth2RegisteredClientAuthConfig : oauth2RegisteredClientAuthorizationConfigs) {
      Oauth2AceAuthorizationCallbackApi oauth2AceAuthorizationCallbackApi = this.callbackApiFactory.getOrCreateUserInfoCallbackApi(oauth2RegisteredClientAuthConfig);
      Oauth2AceUserInfoResultDto newOauth2AceUserInfoResultDto = oauth2AceAuthorizationCallbackApi.getUserInfo(userInfoCallbackInput);
      oauth2AceUserInfoResultDto = this.checkAndMergeUserInfoResult(oauth2AceUserInfoResultDto, newOauth2AceUserInfoResultDto);
    }

    Oauth2AceUserInfoDto oauth2AceUserInfoDto = oauth2AceUserInfoResultDto.check();

    Oauth2UserInfoDto userInfoResultDetailDto = Oauth2UserInfoDto.builder()
        .id(oauth2AceUserInfoDto.getId())
        .attributes(oauth2AceUserInfoDto.getAttributes())
        .authorities(oauth2AceUserInfoDto.getAuthorities())
        .build();

    tokenClaims.claim(Oauth2Constant.TOKEN_CLAIMS_USERINFO_ATTRIBUTE_KEY, userInfoResultDetailDto);

    return tokenClaims;
  }

  private Oauth2AceUserInfoResultDto checkAndMergeUserInfoResult(Oauth2AceUserInfoResultDto oauth2AceUserInfoResultDto, Oauth2AceUserInfoResultDto newOauth2AceUserInfoResultDto) {
    Oauth2AceUserInfoDto newOauth2AceUserInfoDto = newOauth2AceUserInfoResultDto
        .checkAndDataIsNonNull(() -> new RuntimeException("授权失败"));

    if (oauth2AceUserInfoResultDto != null && !StringUtils.equals(oauth2AceUserInfoResultDto.getData().getId(), newOauth2AceUserInfoDto.getId())) {
      throw new BusinessException("授权失败,非常同一账户id");
    }

    Oauth2AceUserInfoResultDto result = this.mergeUserInfoResult(oauth2AceUserInfoResultDto, newOauth2AceUserInfoResultDto);

    return result;
  }

  private Oauth2AceUserInfoResultDto mergeUserInfoResult(Oauth2AceUserInfoResultDto oauth2AceUserInfoResultDto, Oauth2AceUserInfoResultDto newOauth2AceUserInfoResultDto) {
    Integer code = null;
    String message = null;
    String id = null;
    Map<String, Object> extras = new HashMap<>();
    Map<String, String> attributes = new HashMap<>();
    Set<String> authorities = new HashSet<>();

    code = Optional.ofNullable(oauth2AceUserInfoResultDto)
        .map(Oauth2AceUserInfoResultDto::getCode)
        .orElse(null);

    message = Optional.ofNullable(oauth2AceUserInfoResultDto)
        .map(Oauth2AceUserInfoResultDto::getMessage)
        .orElse(null);

    extras.putAll(
        Optional.ofNullable(oauth2AceUserInfoResultDto)
            .map(Oauth2AceUserInfoResultDto::getExtras)
            .orElse(new HashMap<>())
    );

    id = Optional.ofNullable(oauth2AceUserInfoResultDto)
        .map(Oauth2AceUserInfoResultDto::getData)
        .map(Oauth2AceUserInfoDto::getId)
        .orElse(null);

    attributes.putAll(
        Optional.ofNullable(oauth2AceUserInfoResultDto)
            .map(Oauth2AceUserInfoResultDto::getData)
            .map(Oauth2AceUserInfoDto::getAttributes)
            .orElse(new HashMap<>())
    );

    authorities.addAll(
        Optional.ofNullable(oauth2AceUserInfoResultDto)
            .map(Oauth2AceUserInfoResultDto::getData)
            .map(Oauth2AceUserInfoDto::getAuthorities)
            .orElse(new HashSet<>())
    );

    code = Optional.ofNullable(newOauth2AceUserInfoResultDto)
        .map(Oauth2AceUserInfoResultDto::getCode)
        .orElse(code);

    message = Optional.ofNullable(newOauth2AceUserInfoResultDto)
        .map(Oauth2AceUserInfoResultDto::getMessage)
        .orElse(message);

    extras.putAll(
        Optional.ofNullable(newOauth2AceUserInfoResultDto)
            .map(Oauth2AceUserInfoResultDto::getExtras)
            .orElse(new HashMap<>())
    );

    id = Optional.ofNullable(newOauth2AceUserInfoResultDto)
        .map(Oauth2AceUserInfoResultDto::getData)
        .map(Oauth2AceUserInfoDto::getId)
        .orElse(id);

    attributes.putAll(
        Optional.ofNullable(newOauth2AceUserInfoResultDto)
            .map(Oauth2AceUserInfoResultDto::getData)
            .map(Oauth2AceUserInfoDto::getAttributes)
            .orElse(new HashMap<>())
    );

    authorities.addAll(
        Optional.ofNullable(newOauth2AceUserInfoResultDto)
            .map(Oauth2AceUserInfoResultDto::getData)
            .map(Oauth2AceUserInfoDto::getAuthorities)
            .orElse(new HashSet<>())
    );

    Oauth2AceUserInfoResultDto result = new Oauth2AceUserInfoResultDto();
    result.setCode(code);
    result.setMessage(message);
    result.setExtras(extras);
    result.setData(
        Oauth2AceUserInfoDto.builder()
            .id(id)
            .attributes(attributes)
            .authorities(authorities)
            .build()
    );

    return result;
  }

  private boolean containUserInfoScope(Set<String> authorizedScopes) {
    return Optional.ofNullable(authorizedScopes)
        .stream()
        .flatMap(Set::stream)
        .anyMatch(p -> StringUtils.equals(p, Oauth2Constant.CLIENT_SCOPE_USERINFO_KEY));
  }

  private Map<String, Object> convertClaimsIfNecessary(Map<String, Object> claims) {
    Map<String, Object> convertedClaims = new HashMap<>(claims);

    Object value = claims.get(OAuth2TokenIntrospectionClaimNames.ISS);
    if (value != null && !(value instanceof URL)) {
      URL convertedValue = ClaimConversionService.getSharedInstance().convert(value, URL.class);
      if (convertedValue != null) {
        convertedClaims.put(OAuth2TokenIntrospectionClaimNames.ISS, convertedValue);
      }
    }

    value = claims.get(OAuth2TokenIntrospectionClaimNames.SCOPE);
    if (value != null && !(value instanceof List)) {
      Object convertedValue =
          ClaimConversionService.getSharedInstance()
              .convert(value, OBJECT_TYPE_DESCRIPTOR, LIST_STRING_TYPE_DESCRIPTOR);
      if (convertedValue != null) {
        convertedClaims.put(OAuth2TokenIntrospectionClaimNames.SCOPE, convertedValue);
      }
    }

    value = claims.get(OAuth2TokenIntrospectionClaimNames.AUD);
    if (value != null && !(value instanceof List)) {
      Object convertedValue =
          ClaimConversionService.getSharedInstance()
              .convert(value, OBJECT_TYPE_DESCRIPTOR, LIST_STRING_TYPE_DESCRIPTOR);
      if (convertedValue != null) {
        convertedClaims.put(OAuth2TokenIntrospectionClaimNames.AUD, convertedValue);
      }
    }

    return convertedClaims;
  }

  private OAuth2ClientAuthenticationToken getAuthenticatedClientElseThrowInvalidClient(
      Authentication authentication) {
    OAuth2ClientAuthenticationToken clientPrincipal = null;
    if (OAuth2ClientAuthenticationToken.class.isAssignableFrom(
        authentication.getPrincipal().getClass())) {
      clientPrincipal = (OAuth2ClientAuthenticationToken) authentication.getPrincipal();
    }
    if (clientPrincipal != null && clientPrincipal.isAuthenticated()) {
      return clientPrincipal;
    }
    throw new OAuth2AuthenticationException(OAuth2ErrorCodes.INVALID_CLIENT);
  }
}
